namespace NServiceBus.AcceptanceTesting;

using System;
using System.Collections.Concurrent;
using System.Collections.Generic;
using System.Linq;
using System.Threading;
using System.Threading.Tasks;
using Extensibility;
using Unicast.Subscriptions;
using Unicast.Subscriptions.MessageDrivenSubscriptions;

class AcceptanceTestingSubscriptionStorage : ISubscriptionStorage
{
    public Task Subscribe(Subscriber subscriber, MessageType messageType, ContextBag context, CancellationToken cancellationToken = default)
    {
        var dict = storage.GetOrAdd(messageType, type => new ConcurrentDictionary<string, Subscriber>(StringComparer.OrdinalIgnoreCase));

        dict.AddOrUpdate(subscriber.TransportAddress, _ => subscriber, (_, __) => subscriber);
        return Task.CompletedTask;
    }

    public Task Unsubscribe(Subscriber subscriber, MessageType messageType, ContextBag context, CancellationToken cancellationToken = default)
    {
        if (storage.TryGetValue(messageType, out var dict))
        {
            dict.TryRemove(subscriber.TransportAddress, out var _);
        }
        return Task.CompletedTask;
    }

    public Task<IEnumerable<Subscriber>> GetSubscriberAddressesForMessage(IEnumerable<MessageType> messageTypes, ContextBag context, CancellationToken cancellationToken = default)
    {
        var subscribers = messageTypes
            .SelectMany(msgType => storage.TryGetValue(msgType, out var subs) ? subs.Values : [])
            .GroupBy(s => new { s.TransportAddress, s.Endpoint }) // Subscriber does not implement IEquatable<T>
            .Select(g => g.First());

        return Task.FromResult(subscribers);
    }

    readonly ConcurrentDictionary<MessageType, ConcurrentDictionary<string, Subscriber>> storage = new();
}